# 引入3个可训练的权重矩阵Wq,Wk,Wv，一步步实现自注意力机制。
# 这3个矩阵用于将嵌入的输入词元X(i)分别映射为查询向量、键向量和值向量


import torch  # PyTouch 深度学习框架

inputs = torch.tensor(
    [[0.43, 0.15, 0.89],  # Your     (x^1)
     [0.55, 0.87, 0.66],  # journey  (x^2)
     [0.57, 0.85, 0.64],  # starts   (x^3)
     [0.22, 0.58, 0.33],  # with     (x^4)
     [0.77, 0.25, 0.10],  # one      (x^5)
     [0.05, 0.80, 0.55]]  # step     (x^6)
)

x_2 = inputs[1]
d_in = inputs.shape[1]  # 输入嵌入维度 d_in=3
d_out = 2  # 输出嵌入维度 d_out=2
# 在类GPT模型中，输入和输出的维度通常是相同的额，但为了便于理解计算过程，这里使用不同的输入维度d_in=3和输出维度d_out=2


torch.manual_seed(123)
W_query = torch.nn.Parameter(torch.rand(d_in, d_out), requires_grad=False)
W_key = torch.nn.Parameter(torch.rand(d_in, d_out), requires_grad=False)
W_value = torch.nn.Parameter(torch.rand(d_in, d_out), requires_grad=False)
# 设置requires_grad=False以减少输出中的其他项，但如果要在模型训练中使用这些权重矩阵，就需要设置requires_grad=True，以便在训练中更新这些矩阵。

# 接下来，计算查询向量、键向量和值向量
query_2 = x_2 @ W_query
key_2 = x_2 @ W_key
value_2 = x_2 @ W_value
print(query_2)
# tensor([0.4306, 1.4551])


# 权重参数与注意力权重
#  在权重矩阵W中，“权重”是“权重参数”的简称，表示在训练过程中优化的神经网络参数。
#  这与注意力权重是不同的。正如我们已经看到的，注意力权重决定了上下文向量对输入的不同部分的依赖程度（网络对输入的不同部分的关注程度）。
#  总之，权重参数是定义网络连接的基本学习系数，而注意力权重是动态且特定于上下文的值。

keys = inputs @ W_key
values = inputs @ W_value
print("keys.shape: ", keys.shape)
print("values.shape: ", values.shape)
# keys.shape:  torch.Size([6, 2])
# values.shape:  torch.Size([6, 2])


keys_2 = keys[1]
attn_score_22 = query_2.dot(keys_2)
print(attn_score_22)
# tensor(1.8524)

attn_score_2 = query_2 @ keys.T
print(attn_score_2)
# tensor([1.2705, 1.8524, 1.8111, 1.0795, 0.5577, 1.5440])


d_k = keys.shape[-1]
attn_weight_2 = torch.softmax(attn_score_2 / d_k ** 0.5, dim=-1)
print(attn_weight_2)
# tensor([0.1500, 0.2264, 0.2199, 0.1311, 0.0906, 0.1820])


context_vec_2 = attn_weight_2 @ values
print(context_vec_2)
# tensor([0.3061, 0.8210])


